# -*- coding: UTF-8 -*-

import os.path
from gensim.models import ldamodel
from gensim import similarities
import preprocess


class Topic_Model():

    def __init__(self, model_name, docs, corpus, id2word, num_topics=250, model_type='lda'):
        model_name = 'topic_count' + str(num_topics) + '_' + model_name
        self.docs = docs
        self.corpus = corpus
        self.dict = id2word
        if model_type == 'lda':
            if os.path.isfile(model_name):
                self.model = ldamodel.LdaModel.load(model_name, mmap='r')
            else:
                self.model = ldamodel.LdaModel(corpus=corpus, id2word=id2word, num_topics=num_topics)
                self.model.save(model_name)
            self.get_matrix_similarity()
    
    def get_matrix_similarity(self):
        self.index = similarities.MatrixSimilarity(self.model[self.corpus], num_features=len(self.dict))
    
    def get_topn_similar_doc(self, query, topn=10):
        res = []
        query_words = preprocess.sentence_preprocess(query)
        query_bow = self.dict.doc2bow(query_words)
        query_lda = self.model[query_bow]
        # print(query_lda)

        sims = self.index[query_lda]
        sort_sims = sorted(enumerate(sims), key=lambda item: -item[1])
        topnsims = sort_sims[:topn]
        for i in range(len(topnsims)):
            res.append(topnsims[i][0])
            # print(topnsims[i], i, self.docs[topnsims[i][0]])
        return res

    